{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "9d46f78f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "class CliffWalkingEnv:\n",
    "    \"\"\" 悬崖漫步环境\"\"\"\n",
    "    def __init__(self, ncol=12, nrow=4):\n",
    "        # 定义网格世界的列\n",
    "        self.ncol = ncol \n",
    "        # 定义网格世界的行\n",
    "        self.nrow = nrow\n",
    "        # 转移矩阵P[state][action] = [(p, next_state, reward, done)]包含下一个状态和奖励\n",
    "        self.P = self.createP()\n",
    "\n",
    "    # createP 方法用于生成环境的状态转移矩阵 P，该矩阵描述了在每个状态下采取\n",
    "    # 每个动作后可能的转移情况，包括转移概率、下一个状态、奖励和是否终止的信息。\n",
    "    def createP(self):\n",
    "        # 初始化一个三维列表 P，用于存储状态转移矩阵。P 的维度为 (状态数, 动作数, 转移信息)，\n",
    "        # 其中状态数为 self.nrow * self.ncol，动作数为 4。\n",
    "        P = [[[] for j in range(4)] for i in range(self.nrow * self.ncol)]\n",
    "        # 4种动作, change[0]:上,c hange[1]:下, change[2]:左, change[3]:右。坐标系原点(0,0)\n",
    "        # 定义在左上角\n",
    "        # 定义 4 种动作对应的坐标变化，分别表示向上、向下、向左、向右移动。\n",
    "        change = [[0, -1], [0, 1], [-1, 0], [1, 0]]\n",
    "        # 使用三层嵌套循环遍历所有状态和动作。\n",
    "        for i in range(self.nrow):\n",
    "            for j in range(self.ncol):\n",
    "                for a in range(4):\n",
    "                    # 位置在悬崖或者目标状态,因为无法继续交互,任何动作奖励都为0\n",
    "                    # 如果在悬崖或目标状态，无论采取什么动作，都保持在当前状态，\n",
    "                    # 转移概率为 1，奖励为 0，且标记为终止状态\n",
    "                    if i == self.nrow - 1 and j > 0: \n",
    "                        # 判断当前状态是否在悬崖或目标状态（最后一行除第一个位置外）。\n",
    "                        # (3, 1) (3, 2) ... (3, 10)为悬崖 (3, 11)为终点\n",
    "                        # P[3 * 12 + 1][0-3] = [(1, 37, 0, True)]\n",
    "                        # 将转移信息（转移概率、下一个状态、奖励、是否终止）存储到状态转移矩阵 P 中\n",
    "                        P[i * self.ncol +j][a] = [(1, i * self.ncol + j, 0, True)]\n",
    "                        continue\n",
    "                    # 其他位置\n",
    "                    # 计算采取动作 a 后下一个状态的坐标，确保坐标在网格世界范围内。\n",
    "                    '''\n",
    "                    next_x = min(3, max(0, 0 + change[0][1]))\n",
    "                    next_y = min(11, max(0, 0 + change[0][0]))\n",
    "                    '''\n",
    "                    next_x = min(self.nrow - 1, max(0, i + change[a][1])) # 行坐标 大于0 小于 3\n",
    "                    next_y = min(self.ncol - 1, max(0, j + change[a][0])) # 列坐标 大于0 小于 11\n",
    "                    # 将下一个状态的坐标转换为状态编号。\n",
    "                    next_state = next_x * self.ncol + next_y\n",
    "                    # 默认奖励为 -1，表示每走一步都有一个小的惩罚。\n",
    "                    reward = -1\n",
    "                    # 默认不是终止状态。\n",
    "                    done = False\n",
    "                    # 下一个位置在悬崖或者终点\n",
    "                    # 判断下一个状态是否在悬崖或终点。\n",
    "                    if next_x == self.nrow - 1 and next_y > 0:\n",
    "                        done = True\n",
    "                        # 下一个位置在悬崖\n",
    "                        # 判断下一个状态是否在悬崖（除最后一个位置外）。\n",
    "                        if next_y != self.ncol - 1:\n",
    "                            reward = -100\n",
    "                    # 将转移信息（转移概率、下一个状态、奖励、是否终止）存储到状态转移矩阵 P 中。\n",
    "                    P[i * self.ncol + j][a] = [(1, next_state, reward, done)]\n",
    "        return P"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "25f0c295",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PolicyIteration:\n",
    "    \"\"\" 策略迭代算法 \"\"\"\n",
    "    def __init__(self, env, theta, gamma):\n",
    "        # self.env：表示强化学习的环境，该环境需要有 ncol（环境的列数）、nrow（环境的行数）和 P（状态转移概率矩阵）等属性。\n",
    "        self.env = env\n",
    "        # 初始化价值为0\n",
    "        self.v = [0] * self.env.ncol * self.env.nrow\n",
    "        # 初始化为均匀随机策略\n",
    "        self.pi = [[0.25, 0.25, 0.25, 0.25] for i in range(self.env.ncol * self.env.nrow)]\n",
    "        # 策略评估收敛阈值\n",
    "        self.theta = theta\n",
    "        # 折扣因子\n",
    "        self.gamma = gamma\n",
    "\n",
    "    '''\n",
    "    policy_evaluation 方法用于评估当前策略下的状态价值函数。\n",
    "        cnt：用于记录策略评估的迭代轮数。\n",
    "        使用 while 循环进行迭代，直到满足收敛条件（max_diff < self.theta）。\n",
    "        对于每个状态 s，计算所有动作的动作价值函数 Q(s, a)，并根据当前策略 self.pi 对动作价值进行加权求和，得到状态价值函数 V(s)。\n",
    "        max_diff 用于记录每次迭代中状态价值函数的最大变化量，当最大变化量小于阈值 self.theta 时，认为策略评估收敛。\n",
    "    '''\n",
    "    def policy_evaluation(self): \n",
    "        # 计数器\n",
    "        cnt = 1\n",
    "        while True:\n",
    "            max_diff = 0\n",
    "            new_v = [0] * self.env.ncol * self.env.nrow\n",
    "            for s in range(self.env.ncol * self.env.nrow):\n",
    "                # 开始计算状态s下的所有Q(s,a)价值\n",
    "                qsa_list = []\n",
    "                for a in range(4):\n",
    "                    qsa = 0\n",
    "                    for res in self.env.P[s][a]:\n",
    "                        p, next_state, r, done = res\n",
    "                        qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))\n",
    "                        # 本章环境比较特殊,奖励和下一个状态有关,所以需要和状态转移概率相乘\n",
    "                    qsa_list.append(self.pi[s][a] * qsa)\n",
    "                # 状态价值函数和动作价值函数之间的关系\n",
    "                new_v[s] = sum(qsa_list)\n",
    "                max_diff = max(max_diff, abs(new_v[s] - self.v[s]))\n",
    "            self.v = new_v\n",
    "            '''\n",
    "            print(new_v)\n",
    "            [-1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -25.75, -25.75, -25.75, -25.75, -25.75, -25.75, -25.75, -25.75, -25.75, -25.75, -1.0, -25.75, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]\n",
    "            [-1.9, -1.9, -1.9, -1.9, -1.9, -1.9, -1.9, -1.9, -1.9, -1.9, -1.9, -1.9, -1.9, -7.468749999999999, -7.468749999999999, -7.468749999999999, -7.468749999999999, -7.468749999999999, -7.468749999999999, -7.468749999999999, -7.468749999999999, -7.468749999999999, -7.468749999999999, -1.9, -13.0375, -31.993750000000002, -37.5625, -37.5625, -37.5625, -37.5625, -37.5625, -37.5625, -37.5625, -37.5625, -31.993750000000002, -7.2437499999999995, -37.5625, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]\n",
    "            [-2.71, -3.96296875, -3.96296875, -3.96296875, -3.96296875, -3.96296875, -3.96296875, -3.96296875, -3.96296875, -3.96296875, -3.96296875, -2.71, -6.46890625, -10.7340625, -13.239999999999998, -13.239999999999998, -13.239999999999998, -13.239999999999998, -13.239999999999998, -13.239999999999998, -13.239999999999998, -13.239999999999998, -10.7340625, -5.1653125, -20.01109375, -38.81546875, -43.080625000000005, -44.33359375, -44.33359375, -44.33359375, -44.33359375, -44.33359375, -44.33359375, -43.080625, -37.511874999999996, -10.2559375, -45.5865625, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]\n",
    "            [-10.008749065266038, -10.008749065266038, -10.008749065266038, -10.008749065266038, -10.008749065266038, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.008749065266038, -10.008749065266038, -10.008749065266038, -10.008749065266038, -10.008749065266038, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.008749065266038, -10.008749065266038, -10.008749065266038, -10.008749065266038, -10.008749065266038, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.005812617174692, -10.005812617174692, -1.9, -1.0, -10.008749065266038, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]\n",
    "            '''\n",
    "            # 满足收敛条件,退出评估迭代\n",
    "            if max_diff < self.theta: break\n",
    "            cnt += 1\n",
    "        print(\"策略评估进行%d轮后完成\" % cnt)\n",
    "    \n",
    "    '''\n",
    "    policy_improvement 方法用于提升当前策略。\n",
    "        对于每个状态 s，计算所有动作的动作价值函数 Q(s, a)。\n",
    "        找出最大的动作价值 maxq，并统计有多少个动作的动作价值等于 maxq。\n",
    "        将这些动作的概率设置为 1 / cntq，其他动作的概率设置为 0，从而得到新的策略。\n",
    "    '''\n",
    "    # 策略提升\n",
    "    def policy_improvement(self):\n",
    "        for s in range(self.env.nrow * self.env.ncol):\n",
    "            qsa_list = []\n",
    "            for a in range(4):\n",
    "                qsa = 0\n",
    "                for res in self.env.P[s][a]:\n",
    "                    p, next_state, r, done = res\n",
    "                    qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))\n",
    "                qsa_list.append(qsa)\n",
    "            maxq = max(qsa_list)\n",
    "            # 计算有几个动作得到了最大的Q值\n",
    "            cntq = qsa_list.count(maxq)\n",
    "            # 让这些动作均分概率\n",
    "            self.pi[s] = [1 / cntq if q == maxq else 0 for q in qsa_list]\n",
    "        print(\"策略提升完成\")\n",
    "        '''\n",
    "        print(self.pi)\n",
    "        [[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]]\n",
    "        策略提升完成\n",
    "        [[0.5, 0, 0.5, 0], [0, 0, 1.0, 0], [0, 0, 1.0, 0], [0, 0, 1.0, 0], [0, 0, 1.0, 0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0.5, 0, 0, 0.5], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [0, 0, 0, 1.0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [1.0, 0, 0, 0], [0, 0, 0, 1.0], [0, 1.0, 0, 0], [1.0, 0, 0, 0], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]]\n",
    "        策略评估进行72轮后完成\n",
    "        策略提升完成\n",
    "        [[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0, 0, 0, 1.0], [0.3333333333333333, 0.3333333333333333, 0, 0.3333333333333333], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0, 0, 0, 1.0], [0.3333333333333333, 0.3333333333333333, 0, 0.3333333333333333], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0, 1.0, 0, 0], [0, 1.0, 0, 0], [0.25, 0.25, 0.25, 0.25], [0.3333333333333333, 0, 0.3333333333333333, 0.3333333333333333], [0.3333333333333333, 0, 0.3333333333333333, 0.3333333333333333], [0.3333333333333333, 0, 0.3333333333333333, 0.3333333333333333], [0, 0, 0, 1.0], [0.5, 0, 0, 0.5], [0.3333333333333333, 0, 0.3333333333333333, 0.3333333333333333], [0.3333333333333333, 0, 0.3333333333333333, 0.3333333333333333], [0.3333333333333333, 0, 0.3333333333333333, 0.3333333333333333], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 1.0, 0, 0], [0.3333333333333333, 0.3333333333333333, 0.3333333333333333, 0], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]]\n",
    "        策略评估进行44轮后完成\n",
    "        策略提升完成\n",
    "        [[0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 1.0, 0, 0], [0, 1.0, 0, 0], [0, 1.0, 0, 0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 1.0, 0, 0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 1.0, 0, 0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 1.0, 0, 0], [1.0, 0, 0, 0], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]]\n",
    "        策略评估进行12轮后完成\n",
    "        策略提升完成\n",
    "        [[0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 1.0, 0, 0], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 1.0, 0, 0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 1.0, 0, 0], [1.0, 0, 0, 0], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]]\n",
    "        策略评估进行1轮后完成\n",
    "        策略提升完成\n",
    "        [[0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 1.0, 0, 0], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 0.5, 0, 0.5], [0, 1.0, 0, 0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 0, 0, 1.0], [0, 1.0, 0, 0], [1.0, 0, 0, 0], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]]\n",
    "        状态价值：\n",
    "        '''\n",
    "        return self.pi\n",
    "    '''\n",
    "    policy_iteration 方法是策略迭代算法的核心，它通过交替调用 policy_evaluation 和 policy_improvement 方法，不断优化策略。\n",
    "        使用 while 循环进行迭代，直到策略收敛（即新策略和旧策略相同）。\n",
    "        在每次迭代中，先进行策略评估，然后进行策略提升，最后比较新策略和旧策略是否相同。\n",
    "    '''\n",
    "    # 策略迭代\n",
    "    def policy_iteration(self):\n",
    "        while True:\n",
    "            self.policy_evaluation()\n",
    "            old_pi = copy.deepcopy(self.pi)\n",
    "            new_pi = self.policy_improvement()\n",
    "            if old_pi == new_pi: break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e7b04f07",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_agent(agent, action_meaning, disaster=[], end=[]):\n",
    "    print(\"状态价值：\")\n",
    "    for i in range(agent.env.nrow):\n",
    "        for j in range(agent.env.ncol):\n",
    "            # 为了输出美观,保持输出6个字符\n",
    "            print('%6.6s' % ('%.3f' % agent.v[i * agent.env.ncol + j]), end=' ')\n",
    "        print()\n",
    "    \n",
    "    '''\n",
    "    先打印 \"策略：\" 作为提示信息。\n",
    "    同样利用两层嵌套的 for 循环遍历环境中的每个状态。\n",
    "    针对每个状态，存在以下三种处理情况：\n",
    "        若该状态的索引处于 disaster 列表中，就打印 **** 来表示这是灾难状态。\n",
    "        若该状态的索引处于 end 列表中，就打印 EEEE 来表示这是目标状态。\n",
    "        若该状态不属于上述两种特殊情况，就获取该状态下的策略 a。接着遍历 action_meaning 列表，\n",
    "    若策略中对应动作的概率大于 0，就将该动作的含义添加到 pi_str 中；反之，则添加 o。最后打印 pi_str。\n",
    "        内层循环结束后，使用 print() 换行，进入下一行的输出。\n",
    "    '''\n",
    "    print(\"策略：\")\n",
    "    for i in range(agent.env.nrow):\n",
    "        for j in range(agent.env.ncol):\n",
    "            # 一些特殊的状态,例如悬崖漫步中的悬崖\n",
    "            if (i * agent.env.ncol + j) in disaster:\n",
    "                print('****', end=' ')\n",
    "            # 目标状态\n",
    "            elif (i * agent.env.ncol + j) in end:\n",
    "                print('EEEE', end=' ')\n",
    "            else:\n",
    "                a = agent.pi[i * agent.env.ncol + j]\n",
    "                pi_str = ''\n",
    "                for k in range(len(action_meaning)):\n",
    "                    pi_str += action_meaning[k] if a[k] > 0 else 'o'\n",
    "                print(pi_str, end=' ')\n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9e3c0643",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "策略评估进行60轮后完成\n",
      "策略提升完成\n",
      "策略评估进行72轮后完成\n",
      "策略提升完成\n",
      "策略评估进行44轮后完成\n",
      "策略提升完成\n",
      "策略评估进行12轮后完成\n",
      "策略提升完成\n",
      "策略评估进行1轮后完成\n",
      "策略提升完成\n",
      "状态价值：\n",
      "-7.712 -7.458 -7.176 -6.862 -6.513 -6.126 -5.695 -5.217 -4.686 -4.095 -3.439 -2.710 \n",
      "-7.458 -7.176 -6.862 -6.513 -6.126 -5.695 -5.217 -4.686 -4.095 -3.439 -2.710 -1.900 \n",
      "-7.176 -6.862 -6.513 -6.126 -5.695 -5.217 -4.686 -4.095 -3.439 -2.710 -1.900 -1.000 \n",
      "-7.458  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000 \n",
      "策略：\n",
      "ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovoo \n",
      "ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovoo \n",
      "ooo> ooo> ooo> ooo> ooo> ooo> ooo> ooo> ooo> ooo> ooo> ovoo \n",
      "^ooo **** **** **** **** **** **** **** **** **** **** EEEE \n"
     ]
    }
   ],
   "source": [
    "env = CliffWalkingEnv()\n",
    "ep = env.P\n",
    "action_meaning = ['^', 'v', '<', '>']\n",
    "theta = 0.001\n",
    "gamma = 0.9\n",
    "agent = PolicyIteration(env, theta, gamma)\n",
    "agent.policy_iteration()\n",
    "a_pi = agent.pi\n",
    "print_agent(agent, action_meaning, list(range(37, 47)), [47])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "825b370c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ValueIteration:\n",
    "    '''\n",
    "    env：代表强化学习环境。\n",
    "    v：一个列表，用于存储每个状态的价值，初始值都设为 0。\n",
    "    theta：价值收敛的阈值，当两次迭代间最大价值变化小于此值时，迭代终止。\n",
    "    gamma：折扣因子，用于权衡未来奖励的重要性。\n",
    "    pi：一个列表，用来存储每个状态下的策略，初始值为None。\n",
    "    '''\n",
    "    # 价值迭代算法\n",
    "    def __init__(self, env, theta, gamma):\n",
    "        self.env = env\n",
    "        # 初始化价值为0\n",
    "        self.v = [0] * self.env.ncol * self.env.nrow\n",
    "        # 价值收敛阈值\n",
    "        self.theta = theta\n",
    "        self.gamma = gamma\n",
    "        # 价值迭代结束后得到的策略\n",
    "        self.pi = [None for i in range(self.env.ncol * self.env.nrow)]\n",
    "    \n",
    "    '''\n",
    "    cnt：记录迭代的轮数。\n",
    "    max_diff：用于记录两次迭代间状态价值的最大变化。\n",
    "    new_v：存储每次迭代后的新状态价值。\n",
    "    内层循环：\n",
    "        针对每个状态 s，计算所有动作 a 的动作价值 Q(s, a)。\n",
    "        qsa_list 存储状态 s 下所有动作的 Q 值。\n",
    "        env.P[s][a] 包含了在状态 s 执行动作 a 后的所有可能结果，每个结果有概率 p、下一个状态 next_state、奖励 r 和是否结束标志 done。\n",
    "        每个 Q(s, a) 是所有可能结果的期望奖励与下一个状态价值的加权和。\n",
    "    对于每个状态 s，将新的状态价值 new_v[s] 设为 Q(s, a) 的最大值。\n",
    "    当 max_diff 小于 theta 时，迭代结束，调用 get_policy 方法得到最优策略。\n",
    "    '''\n",
    "    def value_iteration(self):\n",
    "        cnt = 0\n",
    "        while True:\n",
    "            max_diff = 0\n",
    "            new_v = [0] * self.env.ncol * self.env.nrow\n",
    "            for s in range(self.env.ncol * self.env.nrow):\n",
    "                # 开始计算状态s下的所有Q(s,a)价值\n",
    "                qsa_list = []\n",
    "                for a in range(4):\n",
    "                    qsa = 0\n",
    "                    for res in self.env.P[s][a]:\n",
    "                        p, next_state, r, done = res\n",
    "                        qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))\n",
    "                    # 这一行和下一行代码是价值迭代和策略迭代的主要区别\n",
    "                    qsa_list.append(qsa)\n",
    "                new_v[s] = max(qsa_list)\n",
    "                max_diff = max(max_diff, abs(new_v[s] - self.v[s]))\n",
    "            self.v = new_v\n",
    "            # 满足收敛条件,退出评估迭代\n",
    "            if max_diff < self.theta: break\n",
    "            cnt += 1\n",
    "        print(\"价值迭代一共进行%d轮\" % cnt)\n",
    "        self.get_policy()\n",
    "\n",
    "    '''\n",
    "    针对每个状态 s，重新计算所有动作的 Q 值。\n",
    "    找出最大的 Q 值 maxq，统计有多少个动作能达到这个最大值 cntq。\n",
    "    让这些能达到最大 Q 值的动作均分概率，其余动作概率为 0，以此生成贪婪策略。\n",
    "    '''\n",
    "    # 根据价值函数导出一个贪婪策略\n",
    "    def get_policy(self):\n",
    "        for s in range(self.env.nrow * self.env.ncol):\n",
    "            qsa_list = []\n",
    "            for a in range(4):\n",
    "                qsa = 0\n",
    "                for res in self.env.P[s][a]:\n",
    "                    p, next_state, r, done = res\n",
    "                    qsa += p * (r + self.gamma * self.v[next_state] * (1 - done))\n",
    "                qsa_list.append(qsa)\n",
    "            maxq = max(qsa_list)\n",
    "            # 计算有几个动作得到了最大的Q值 \n",
    "            cntq = qsa_list.count(maxq)\n",
    "            # 让这些动作均分概率\n",
    "            self.pi[s] = [1 / cntq if q == maxq else 0 for q in qsa_list]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7f46dd2f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "价值迭代一共进行14轮\n",
      "状态价值：\n",
      "-7.712 -7.458 -7.176 -6.862 -6.513 -6.126 -5.695 -5.217 -4.686 -4.095 -3.439 -2.710 \n",
      "-7.458 -7.176 -6.862 -6.513 -6.126 -5.695 -5.217 -4.686 -4.095 -3.439 -2.710 -1.900 \n",
      "-7.176 -6.862 -6.513 -6.126 -5.695 -5.217 -4.686 -4.095 -3.439 -2.710 -1.900 -1.000 \n",
      "-7.458  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000 \n",
      "策略：\n",
      "ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovoo \n",
      "ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovo> ovoo \n",
      "ooo> ooo> ooo> ooo> ooo> ooo> ooo> ooo> ooo> ooo> ooo> ovoo \n",
      "^ooo **** **** **** **** **** **** **** **** **** **** EEEE \n"
     ]
    }
   ],
   "source": [
    "env = CliffWalkingEnv()\n",
    "action_meaning = ['^', 'v', '<', '>']\n",
    "theta = 0.001\n",
    "gamma = 0.9\n",
    "agent = ValueIteration(env, theta, gamma)\n",
    "agent.value_iteration()\n",
    "print_agent(agent, action_meaning, list(range(37, 47)), [47])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96520200",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "冰洞的索引: {11, 12, 5, 7}\n",
      "目标的索引: {15}\n",
      "[(0.3333333333333333, 10, 0.0, False), (0.3333333333333333, 13, 0.0, False), (0.3333333333333333, 14, 0.0, False)]\n",
      "[(0.3333333333333333, 13, 0.0, False), (0.3333333333333333, 14, 0.0, False), (0.3333333333333333, 15, 1.0, True)]\n",
      "[(0.3333333333333333, 14, 0.0, False), (0.3333333333333333, 15, 1.0, True), (0.3333333333333333, 10, 0.0, False)]\n",
      "[(0.3333333333333333, 15, 1.0, True), (0.3333333333333333, 10, 0.0, False), (0.3333333333333333, 13, 0.0, False)]\n"
     ]
    }
   ],
   "source": [
    "import gym\n",
    "# 创建环境\n",
    "env = gym.make(\"FrozenLake-v1\")\n",
    "# 解封装才能访问状态转移矩阵P\n",
    "env = env.unwrapped\n",
    "# 环境渲染,通常是弹窗显示或打印出可视化的环境\n",
    "env.render()\n",
    "\n",
    "holes = set()\n",
    "ends = set()\n",
    "# 这里使用了三重循环来遍历状态转移矩阵 P。外层循环遍历所有的状态 s，中层循环遍历在每个状态下可以采取的所有动作 a，内层循环遍历在状态 s 下采取动作 a 后可能转移到的所有下一个状态 s_。\n",
    "for s in env.P:\n",
    "    for a in env.P[s]:\n",
    "        for s_ in env.P[s][a]:\n",
    "            # 获得奖励为1,代表是目标\n",
    "            # 在状态转移矩阵中，s_ 是一个元组，其中 s_[2] 表示转移到下一个状态后获得的奖励。如果奖励为 1.0，说明这个状态是目标位置，将其添加到 ends 集合中。\n",
    "            if s_[2] == 1.0:\n",
    "                ends.add(s_[1])\n",
    "            # s_[3] 表示转移到下一个状态后是否为终止状态（即是否结束当前回合）。如果为 True，说明这个状态是一个终止状态，将其添加到 holes 集合中。\n",
    "            if s_[3] == True:\n",
    "                holes.add(s_[1])\n",
    "# 最后，从 holes 集合中去除 ends 集合中的元素，确保 holes 集合中只包含洞的位置，不包含目标位置。\n",
    "holes = holes - ends\n",
    "print(\"冰洞的索引:\", holes)\n",
    "print(\"目标的索引:\", ends)\n",
    "\n",
    "# 查看目标左边一格的状态转移信息\n",
    "for a in env.P[14]:\n",
    "    print(env.P[14][a])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "808a0d79",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "策略评估进行25轮后完成\n",
      "策略提升完成\n",
      "策略评估进行58轮后完成\n",
      "策略提升完成\n",
      "状态价值：\n",
      " 0.069  0.061  0.074  0.056 \n",
      " 0.092  0.000  0.112  0.000 \n",
      " 0.145  0.247  0.300  0.000 \n",
      " 0.000  0.380  0.639  0.000 \n",
      "策略：\n",
      "<ooo ooo^ <ooo ooo^ \n",
      "<ooo **** <o>o **** \n",
      "ooo^ ovoo <ooo **** \n",
      "**** oo>o ovoo EEEE \n"
     ]
    }
   ],
   "source": [
    "# 这个动作意义是Gym库针对冰湖环境事先规定好的\n",
    "action_meaning = ['<', 'v', '>', '^']\n",
    "theta = 1e-5\n",
    "gamma = 0.9\n",
    "agent = PolicyIteration(env, theta, gamma)\n",
    "agent.policy_iteration()\n",
    "print_agent(agent, action_meaning, [5, 7, 11, 12], [15])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d4753b54",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "价值迭代一共进行60轮\n",
      "状态价值：\n",
      " 0.069  0.061  0.074  0.056 \n",
      " 0.092  0.000  0.112  0.000 \n",
      " 0.145  0.247  0.300  0.000 \n",
      " 0.000  0.380  0.639  0.000 \n",
      "策略：\n",
      "<ooo ooo^ <ooo ooo^ \n",
      "<ooo **** <o>o **** \n",
      "ooo^ ovoo <ooo **** \n",
      "**** oo>o ovoo EEEE \n"
     ]
    }
   ],
   "source": [
    "action_meaning = ['<', 'v', '>', '^']\n",
    "theta = 1e-5\n",
    "gamma = 0.9\n",
    "agent = ValueIteration(env, theta, gamma)\n",
    "agent.value_iteration()\n",
    "print_agent(agent, action_meaning, [5, 7, 11, 12], [15])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
